import json
import time
import os
import re
import Sunny as SyNet
import threading
import socket
import urllib3
import requests
from queue import Queue
from collections import defaultdict

# 禁用 SSL 警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

def BytesToStr(b):
    """将字节数组转为字符串，尝试多种编码"""
    encodings = ['utf-8', 'gbk', 'gb2312', 'gb18030', 'latin1']
    for encoding in encodings:
        try:
            return b.decode(encoding)
        except UnicodeDecodeError:
            continue
    # 如果所有编码都失败，使用 errors='ignore' 忽略错误字符
    return b.decode('utf-8', errors='ignore')

def clean_filename(filename):
    """清理文件名中的非法字符"""
    return re.sub(r'[\\/*?:"<>|]', '_', filename)

class FolderManager:
    """文件夹管理类"""
    def __init__(self, base_dir):
        self.base_dir = base_dir
        self.current_folder = None
        self.current_count = 0
        self.folder_lock = threading.Lock()
        self._init_folder()

    def _init_folder(self):
        """初始化文件夹信息"""
        os.makedirs(self.base_dir, exist_ok=True)
        with self.folder_lock:
            folders = [f for f in os.listdir(self.base_dir) 
                      if os.path.isdir(os.path.join(self.base_dir, f)) and f.isdigit()]
            
            if not folders:
                self.current_folder = "01"
            else:
                self.current_folder = max(folders)
            
            folder_path = os.path.join(self.base_dir, self.current_folder)
            os.makedirs(folder_path, exist_ok=True)
            self.current_count = len([f for f in os.listdir(folder_path) 
                                    if f.endswith('.mp4')])

    def get_path(self, filename):
        """获取文件保存路径"""
        with self.folder_lock:
            if self.current_count >= 3000:
                # 创建新文件夹
                next_num = int(self.current_folder) + 1
                self.current_folder = f"{next_num:02d}"
                self.current_count = 0
                folder_path = os.path.join(self.base_dir, self.current_folder)
                os.makedirs(folder_path, exist_ok=True)
                print(f"创建新文件夹: {self.current_folder}")
            
            self.current_count += 1
            return os.path.join(self.base_dir, self.current_folder, filename)

class VideoDownloader:
    def __init__(self):
        self.vid_lock = threading.Lock()
        self.stats_lock = threading.Lock()
        self.completed_count = 0
        self.total_count = 0
        self.downloaded_vids = set()
        self.download_queue = Queue()
        
        # 加载已下载记录
        self._load_downloaded_vids()
        
        # 启动单个下载线程
        self.download_thread = threading.Thread(target=self._process_downloads, daemon=True)
        self.download_thread.start()

    def _load_downloaded_vids(self):
        """加载已下载记录"""
        try:
            if os.path.exists("vids.txt"):
                with open("vids.txt", 'r', encoding='utf-8') as f:
                    self.downloaded_vids = set(f.read().splitlines())
        except Exception as e:
            print(f"加载下载记录失败: {str(e)}")

    def _process_downloads(self):
        """处理下载队列"""
        while True:
            try:
                # 从队列获取任务
                url, file_path, vid, vids_file = self.download_queue.get()
                
                # 开始下载
                print(f"开始下载: {os.path.basename(file_path)}")
                temp_file = f"{file_path}.tmp"
                
                # 使用 urllib3 进行下载
                http = urllib3.PoolManager()
                with http.request('GET', url, preload_content=False) as resp, \
                     open(temp_file, 'wb') as out_file:
                    # 使用更大的块大小
                    for chunk in resp.stream(232768):
                        out_file.write(chunk)
                
                # 下载完成后移动文件
                os.replace(temp_file, file_path)
                
                # 记录下载信息
                with self.vid_lock:
                    self.downloaded_vids.add(vid)
                    with open(vids_file, 'a', encoding='utf-8') as f:
                        f.write(f"{vid}\n")
                
                # 更新统计
                with self.stats_lock:
                    self.completed_count += 1
                    remaining = self.total_count - self.completed_count
                    print(f"\n下载完成: {os.path.basename(file_path)}")
                    print(f"进度: {self.completed_count}/{self.total_count} | 剩余: {remaining}")
                
            except Exception as e:
                print(f"下载失败: {str(e)}")
                if os.path.exists(temp_file):
                    try:
                        os.remove(temp_file)
                    except:
                        pass
            finally:
                self.download_queue.task_done()
                time.sleep(0.1)  # 添加小延迟

    def add_download_task(self, url, file_path, vid, vids_file):
        """添加下载任务"""
        if vid in self.downloaded_vids:
            print(f"视频已下载过: {os.path.basename(file_path)}")
            return

        with self.stats_lock:
            self.total_count += 1
            print(f"新增下载任务 (总任务数: {self.total_count})")
        
        self.download_queue.put((url, file_path, vid, vids_file))

class Callback:
    Http_消息类型_发起请求 = 1
    Http_消息类型_请求完成 = 2
    Http_消息类型_请求失败 = 3

    downloader = VideoDownloader()
    folder_manager = None

    @staticmethod
    def process_response(SyHTTP, response_content):
        """处理响应内容"""
        try:
            # 基础保存目录
            base_save_dir = "D:/Personal/Desktop/郁金香剪辑/采集"
            
            # 初始化文件夹管理器
            if Callback.folder_manager is None:
                Callback.folder_manager = FolderManager(base_save_dir)
                Callback.downloader.folder_manager = Callback.folder_manager
            
            # 解析JSON
            json_data = json.loads(response_content)
            if 'data' not in json_data or 'list' not in json_data['data']:
                return
            
            # 处理视频列表
            for item in json_data['data']['list']:
                try:
                    if 'playInfo' not in item or 'videoUrl' not in item['playInfo']:
                        continue
                        
                    title = item.get('title', '')
                    video_url = item['playInfo']['videoUrl']
                    vid = item['playInfo'].get('vid', '')
                    article_sku_id = item.get('articleSkuIds', [''])[0]
                    
                    # 检查商品ID
                    if not article_sku_id or article_sku_id == '0' or len(article_sku_id) < 5:
                        print(f"跳过无效商品ID的视频: {title} (ID: {article_sku_id})")
                        continue
                    
                    # 构建文件名
                    file_name = f"{title}-{article_sku_id}.mp4"
                    clean_file_name = clean_filename(file_name)
                    file_path = Callback.folder_manager.get_path(clean_file_name)
                    
                    # 添加到下载队列
                    print(f"\n添加到下载队列: {clean_file_name}")
                    Callback.downloader.add_download_task(
                        video_url,
                        file_path,
                        vid,
                        "vids.txt"
                    )
                    
                except Exception as e:
                    print(f"处理视频失败: {str(e)}")
                    continue
                    
        except json.JSONDecodeError as e:
            print(f"JSON解析错误: {str(e)}")
        except Exception as e:
            print(f"处理响应失败: {str(e)}")

    def Http(SunnyContext, 请求唯一ID, MessageId, 消息类型, 请求方式, 请求地址, 错误信息, pid):
        """HTTP请求回调处理"""
        SyHTTP = SyNet.MessageIdToSunny(MessageId)
        
        if 消息类型 == Callback.Http_消息类型_请求完成:
            try:
                request_url = BytesToStr(请求地址)
                if 'functionId=getRecommendNewFeedsList' in request_url:
                    try:
                        # 处理压缩
                        Encoding = SyHTTP.响应.取协议头("Content-Encoding").lower()
                        响应内容 = None
                        
                        if Encoding == "gzip":
                            Body = SyHTTP.响应.取响应Body()
                            响应内容 = SyNet.Gzip解压缩(Body)
                        elif Encoding == "br":
                            Body = SyHTTP.响应.取响应Body()
                            响应内容 = SyNet.Br解压缩(Body)
                        elif Encoding == "deflate":
                            Body = SyHTTP.响应.取响应Body()
                            响应内容 = SyNet.Deflate解压缩(Body)
                        elif Encoding == "zstd":
                            Body = SyHTTP.响应.取响应Body()
                            响应内容 = SyNet.ZSTD解压缩(Body)
                        else:
                            响应内容 = SyHTTP.响应.取响应Body()
                        
                        # 确保响应内容是字符串
                        response_text = BytesToStr(响应内容)
                        
                        # 异步处理响应
                        threading.Thread(
                            target=Callback.process_response,
                            args=(SyHTTP, response_text),
                            daemon=True
                        ).start()
                        
                    except Exception as e:
                        print(f"处理请求失败: {str(e)}")
                        # 打印详细错误信息以便调试
                        import traceback
                        traceback.print_exc()
            except Exception as e:
                print(f"处理URL失败: {str(e)}")

def find_available_port(start_port=25801, max_attempts=10):
    """查找可用端口"""
    for port in range(start_port, start_port + max_attempts):
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.bind(('0.0.0.0', port))
            sock.close()
            return port
        except:
            continue
    return None

def main():
    # 创建实例
    Sunny = SyNet.SunnyNet()
    
    try:
        # 查找端口
        port = find_available_port()
        if not port:
            print("无法找到可用端口")
            return
            
        # 配置代理
        Sunny.绑定端口(port)
        try:
            Sunny.安装证书()
        except Exception as e:
            print(f"证书安装失败(不影响使用): {str(e)}")
        
        Sunny.绑定回调地址(Callback.Http, None, None, None)
        
        # 启动代理
        if not Sunny.启动():
            print("启动失败")
            print(Sunny.取错误())
            return
        
        print(f"代理服务器已启动在 0.0.0.0:{port}")
        
        # 加载进程代理
        if not Sunny.进程代理_加载驱动():
            print("加载驱动失败，进程代理不可用(注意，需要管理员权限，win7请安装 KB3033929 补丁)")
        else:
            print("进程代理驱动加载成功")
            Sunny.进程代理_添加进程名("Charles.exe")
            print("已添加进程代理: Charles.exe")
        
        # 运行
        while True:
            time.sleep(1)
            
    except KeyboardInterrupt:
        print("\n正在关闭代理服务器...")
        Sunny.停止代理()
    except Exception as e:
        print(f"发生错误: {str(e)}")
        Sunny.停止代理()

if __name__ == "__main__":
    main() 